# -*- coding: utf-8 -*-
"""
Created on Thu Oct 20 18:39:23 2022

@author: ASUS
"""
import numpy as np
import mindspore as ms
import mindspore.nn as nn

def Dice(predict, label, eps=1e-6):
    inter = np.sum(predict * label)
    union = np.sum(predict) + np.sum(label)
    print("11",np.sum(predict),np.sum(label))
    print("inter", inter,"union",union)
    dice = (2*inter + eps) / (union + eps)
    return dice

def Hec_dice(predict, label):
    """
    0:其他组织 1:肾脏 2:肿瘤
    hec1:肾脏+肿瘤
    hec2:肿瘤
    """
    hec1 = Dice(predict > 0, label > 0)
    hec2 = Dice(predict == 2, label == 2)
    return (hec1 + hec2) / 2

class HEC(nn.Metric):
    def __init__(self):
        super(HEC, self).__init__()
        self.clear()

    def clear(self):
        """初始化变量_abs_error_sum和_samples_num"""
        self._abs_error_sum = 0  # 保存误差和
        self._samples_num = 0    # 累计数据量

    @nn.rearrange_inputs
    def update(self, *inputs):
        """更新_abs_error_sum和_samples_num"""
        y_pred = inputs[0].asnumpy()
        y = inputs[1].asnumpy()
        print(y.shape)

        # 计算预测值与真实值的绝对误差
        hec_dice = Hec_dice(y_pred,y)
        print("hec_dice:",hec_dice)
        self._abs_error_sum += hec_dice

        # 样本的总数
        self._samples_num += y.shape[0]

    def eval(self):
        """计算最终评估结果"""
        return self._abs_error_sum / self._samples_num

#     @nn.rearrange_inputs
#     def update(self, *inputs):
#         """更新_abs_error_sum和_samples_num"""
#         y_pred = inputs[0].asnumpy()
#         y = inputs[1].asnumpy()
#         print("before:",y_pred.shape,y.shape)
#         for i in range(y_pred.shape[0]):
#             y_pred2 = y_pred[i]
#             y2 = y[i]
#             y_pred2 = y_pred2.argmax(axis=0)
#             y2 = y2.argmax(axis=0)
#             # 计算预测值与真实值的绝对误差
#             hec_dice = Hec_dice(y_pred2,y2)
#     #         print("预测值：",y_pred)
#     #         print("真实值：",y)
#             print("hec_dice:",i,hec_dice)
#             self._abs_error_sum += hec_dice
#             if i==0:
#                 print("later:",y_pred2.shape,y2.shape)

#         # 样本的总数
#         self._samples_num += y.shape[0]